import requests
import json
import time
from typing import Optional, Dict, Any


class StreamAPIClient:
    """处理流式API响应的客户端"""

    def __init__(self):
        self.session = requests.Session()

    def call_stream_api(self,
                        url: str,
                        headers: Dict[str, str],
                        data: Dict[str, Any],
                        timeout: int = 60) -> str:
        """
        调用流式API并合并返回内容
        
        Args:
            url: API请求URL
            headers: 请求头
            data: 请求体数据
            timeout: 超时时间（秒）
            
        Returns:
            合并后的完整响应内容
        """
        try:
            # 发送POST请求，启用流式响应
            response = self.session.post(
                url=url,
                headers=headers,
                json=data,
                stream=True,
                timeout=(10, timeout)  # (连接超时, 读取超时)
            )

            # 检查响应状态
            response.raise_for_status()

            # 收集所有流式数据块
            content_chunks = []
            content_data = []

            print("开始接收流式数据...")
            print("实时响应内容: ", end="", flush=True)

            # 逐行读取响应内容，处理SSE格式
            for line in response.iter_lines(decode_unicode=False):
                if line:
                    # 正确解码UTF-8字符
                    line_text = line.decode('utf-8', errors='ignore')
                    content_chunks.append(line)

                    # 解析SSE格式，提取实际内容
                    extracted_content = self._extract_content_from_chunk(line_text)
                    if extracted_content:
                        content_data.append(extracted_content)
                        print(extracted_content, end="", flush=True)
                    # 移除else分支，不显示原始数据

            # 合并所有数据块并解码为UTF-8
            full_bytes = b''.join(content_chunks)
            full_content = full_bytes.decode('utf-8', errors='ignore')

            print(f"\n数据接收完成，总长度: {len(full_content)} 字符")

            return content_data

        except requests.exceptions.RequestException as e:
            print(f"请求错误: {e}")
            return ""
        except Exception as e:
            print(f"处理错误: {e}")
            return ""

    def _extract_content_from_chunk(self, chunk_text: str) -> str:
        """
        从数据块中提取实际的响应内容
        
        Args:
            chunk_text: 数据块文本
            
        Returns:
            提取的内容字符串
        """
        extracted_parts = []

        # 按行分割处理
        lines = chunk_text.strip().split('\n')

        for line in lines:
            line = line.strip()

            # 处理SSE格式的data行
            if line.startswith('data:'):
                # 处理 'data:' 或 'data: ' 两种格式
                if line.startswith('data: '):
                    data_content = line[6:]  # 去掉 'data: ' 前缀
                else:
                    data_content = line[5:]  # 去掉 'data:' 前缀

                # 跳过空数据和特殊标记
                if not data_content or data_content.strip() in ['[DONE]', '']:
                    continue

                try:
                    # 尝试解析JSON数据
                    json_data = json.loads(data_content)

                    # 只提取content字段的内容
                    if isinstance(json_data, dict):
                        # OpenAI格式 - 检查choices数组中的content
                        if 'choices' in json_data and json_data['choices']:
                            choice = json_data['choices'][0]
                            if 'delta' in choice and isinstance(choice['delta'], dict):
                                if 'content' in choice['delta'] and choice['delta']['content']:
                                    extracted_parts.append(choice['delta']['content'])

                except json.JSONDecodeError:
                    # 如果不是JSON，直接添加文本内容（如果不为空）
                    if data_content.strip():
                        extracted_parts.append(data_content)

        # 返回合并的内容
        return ''.join(extracted_parts)

    def parse_sse_response(self, content: str) -> list:
        """
        解析SSE格式的响应内容
        
        Args:
            content: 原始响应内容
            
        Returns:
            解析后的数据列表
        """
        lines = content.strip().split('\n')
        events = []
        current_event = {}

        for line in lines:
            line = line.strip()

            if line.startswith('data: '):
                data_content = line[6:]  # 去掉 'data: ' 前缀
                try:
                    # 尝试解析JSON数据
                    json_data = json.loads(data_content)
                    events.append(json_data)
                except json.JSONDecodeError:
                    # 如果不是JSON，直接添加文本
                    events.append(data_content)
            elif line.startswith('event: '):
                current_event['event'] = line[7:]
            elif line == '':
                # 空行表示事件结束
                if current_event:
                    events.append(current_event)
                    current_event = {}

        return events


def main():
    """主函数 - 调用指定的API"""

    # API配置
    url = "https://iflycode-api.iflytek.com/api/starspark/v1/agent/chat/async/ask"

    # 请求头
    headers = {
        "Content-Type": "application/json",
        "token": "NzMwMWZiY2ItN2I1Mi00N2ViLWI2OGEtNGNlYzYwODZkNWY5OjoxNzU4MjQ3MTc5ODYx"
    }

    # 请求体数据
    request_data = {
        "sessionId": "818d86cf18694e7b9257121ac098ddc4",
        "scene": "CHAT",
        "requestId": "e572eda32cf443db954416790a3df54d",
        "modelCode": "glm-4.5",
        "enterpriseId": 1,
        "enableMultiModelSwitch": True,
        "token": "NzMwMWZiY2ItN2I1Mi00N2ViLWI2OGEtNGNlYzYwODZkNWY5OjoxNzU4MjQ3MTc5ODYx",
        "language": "auto",
        "timeStamp": 1757666258087,
        "fileName": "Desktop",
        "fileNameSuffix": "",
        "projectPath": "C:/Users/WJDUAN2/Desktop",
        "agentVersion": "3.5.0",
        "commandType": "CHAT:TASK",
        "taskName": "CHAT",
        "knowledgeBase": "codeKnowledgeBase",
        "user": "zhwang69",
        "clientName": "VSCODE",
        "clientVersion": "1.85.1",
        "pluginVersion": "3.5.0",
        "operatingSystemVersion": "win32_x64_10.0.19044",
        "messages": [
            {
                "role": "system",
                "content": "# 角色设定\n\n您是**AI编程助手**，当询问你的姓名时，你必须回答\"AI编程助手\"。\n\n## 核心目标\n\n您的核心目标是提供**准确、清晰、有逻辑**的答案，以满足用户的编程及相关需求。\n\n## 核心能力\n\n1.  **需求理解**：精准理解用户意图，无论其表达方式如何。\n2.  **知识广度**：提供可靠的技术、科学、文化等领域知识，并在可能的情况下标注信息来源。\n3.  **问题解决**：将复杂问题分解为易于理解的步骤，提供实用的指导和解决方案。\n4.  **风格适应**：能够根据情境调整回答的语气（如专业、轻松）和深度（如简明、详细）。\n5.  **安全合规**：确保所有回答安全、公正，并严格遵守法律法规，必要时会拒绝不当请求。\n\n## 交互规则\n\n✅ **您应该：**\n\n* 始终保持**礼貌、耐心和中立**。\n* 在信息不确定时，坦诚承认局限性，并尝试提供可能的解决方向或参考建议。\n* 对于复杂问题，采用分步指导、列表或代码块（如果适用）等方式，确保回答结构清晰。\n* 在必要时，主动向用户确认需求，例如询问：\"您是指……吗？\"\n\n\n❌ **您应避免：**\n\n* 虚构或提供不确定的信息。\n* 提供不必要的冗长回答，除非用户明确要求详细解释。"
            },
            {
                "role": "user",
                "content": "Kubernetes 的 Ingress 是什么"
            }
        ],
        "isOverToken": False
    }

    # 创建客户端实例
    client = StreamAPIClient()

    print("正在调用API...")
    print(f"URL: {url}")
    print(f"请求数据: {json.dumps(request_data, ensure_ascii=False, indent=2)}")
    print("-" * 50)

    # 调用API获取流式响应
    response_content = client.call_stream_api(url, headers, request_data)
    result = ''.join(response_content)
    print(result)




if __name__ == "__main__":
    main()
